package net.pragyah.scalby.inference.exact

import junit.framework.TestCase
import junit.framework.Assert._
//import scala.collection.mutable.Map

import net.pragyah.scalby._

import net.pragyah.scalgorithms.graphs._

class TestMessagePassing_SinglyConnected extends TestCase{

	val b = new Variable[String]("B",Set[String]()+"b1"+"b2",None) 
	val f = new Variable[String]("F",Set[String]()+"f1"+"f2",None) 
	val a = new Variable[String]("A",Set[String]()+"a1"+"a2",None) 
	

	val B = Vertex[Variable[String]](b);
	val F = Vertex[Variable[String]](f);
	val A = Vertex[Variable[String]](a);
	
	val P = ProbabilityDistribution[String]()
	
	val graph:Graph[Variable[String]] = Graph[Variable[String]](Nil,true)
 
   def atestInitialTree{
	          
	
	        
	        val algo = new MessagePassing_SinglyConnected[String](graph,P,Set()+B+F)
	        
	        algo.initial_tree;
	        
	        
	        val bTag = B.tag.asInstanceOf[VertexTag[String]]
	        assertEquals(0.005f,bTag.P("b1"))
	        assertEquals(0.995f,bTag.P("b2"))
	        assertEquals(1f,bTag.lambda("b1"))
	        assertEquals(1f,bTag.lambda("b2"))
	        assertEquals(0.005f,bTag.pi("b1"))
	        assertEquals(0.995f,bTag.pi("b2"))
	        assert(b() == None)
         
         	val fTag = F.tag.asInstanceOf[VertexTag[String]]
	        assertEquals(0.03f,fTag.P("f1"))
	        assertEquals(0.97f,fTag.P("f2"))
	        assertEquals(1f,fTag.lambda("f1"))
	        assertEquals(1f,fTag.lambda("f2"))
	        assertEquals(0.03f,fTag.pi("f1"))
	        assertEquals(0.97f,fTag.pi("f2"))
	        assert(f() == None)
          
         	val aTag = A.tag.asInstanceOf[VertexTag[String]]
	        assertEquals(14,Math.round(aTag.P("a1")*1000))
	        assertEquals(986,Math.round(aTag.P("a2")*1000))
	        assertEquals(1f,aTag.lambda("a1"))
	        assertEquals(1f,aTag.lambda("a2"))
	        assertEquals(14,Math.round(aTag.pi("a1")*1000))
	        assertEquals(986,Math.round(aTag.pi("a2")*1000))
	        assert(a() == None)
         
         
         
         

	}


   def atestInitialUpdate_A_a1{
	        val algo = new MessagePassing_SinglyConnected[String](graph,P,Set()+B+F)
	        
	        algo.initial_tree;
           algo.update_tree(a,"a1")

           	        
	        val bTag = B.tag.asInstanceOf[VertexTag[String]]
	        assertEquals(358,Math.round(bTag.P("b1")*1000))
	        assertEquals(642,Math.round(bTag.P("b2")*1000))
	        assertEquals(99,Math.round(bTag.lambda("b1")*100))
	        assertEquals(9,Math.round(bTag.lambda("b2")*1000))
	        assertEquals(0.005f,bTag.pi("b1"))
	        assertEquals(0.995f,bTag.pi("b2"))
	        assert(b() == None)
         
         	val fTag = F.tag.asInstanceOf[VertexTag[String]]
	        assertEquals(443,Math.round(fTag.P("f1")*1000))
	        assertEquals(557,Math.round(fTag.P("f2")*1000))
	        assertEquals(204,Math.round(fTag.lambda("f1")*1000))
	        assertEquals(8,Math.round(fTag.lambda("f2")*1000))
	        assertEquals(0.03f,fTag.pi("f1"))
	        assertEquals(0.97f,fTag.pi("f2"))
	        assert(f() == None)
          
         	val aTag = A.tag.asInstanceOf[VertexTag[String]]
	        assertEquals(1.0f,aTag.P("a1"))
	        assertEquals(0.0f,aTag.P("a2"))
	        assertEquals(1.0f,aTag.lambda("a1"))
	        assertEquals(0.0f,aTag.lambda("a2"))
	        assertEquals(1.0f,aTag.pi("a1"))
	        assertEquals(0.0f,aTag.pi("a2"))
         
	        assert(a().get.equals("a1"))
         
         
         
           
   }

   
   
   def testInitialUpdate_A_a1_F_f1{
		  val algo = new MessagePassing_SinglyConnected[String](graph,P,Set()+B+F)
		    
		  algo.initial_tree;
		  algo.update_tree(a,"a1")
		  algo.update_tree(f,"f1")
   
   	        val bTag = B.tag.asInstanceOf[VertexTag[String]]
	        assertEquals(24,Math.round(bTag.P("b1")*1000))
	        assertEquals(976,Math.round(bTag.P("b2")*1000))
	        assertEquals(992,Math.round(bTag.lambda("b1")*1000))
	        assertEquals(2,Math.round(bTag.lambda("b2")*10))
	        assertEquals(0.005f,bTag.pi("b1"))
	        assertEquals(0.995f,bTag.pi("b2"))
	        assert(b() == None)

         	val fTag = F.tag.asInstanceOf[VertexTag[String]]
	        assertEquals(1.0f,fTag.P("f1"))
	        assertEquals(0.0f,fTag.P("f2"))
	        assertEquals(1.0f,fTag.lambda("f1"))
	        assertEquals(0.0f,fTag.lambda("f2"))
	        assertEquals(1.0f,fTag.pi("f1"))
	        assertEquals(0.0f,fTag.pi("f2"))
         
	        assert(f().get.equals("f1"))
          
         	val aTag = A.tag.asInstanceOf[VertexTag[String]]
	        assertEquals(1.0f,aTag.P("a1"))
	        assertEquals(0.0f,aTag.P("a2"))
	        assertEquals(1.0f,aTag.lambda("a1"))
	        assertEquals(0.0f,aTag.lambda("a2"))
	        assertEquals(1.0f,aTag.pi("a1"))
	        assertEquals(0.0f,aTag.pi("a2"))
         
	        assert(a().get.equals("a1"))

	
   }
   
   //TODO write test cases for each methods
   
    /// SETUP 
    
	    val all = B::F::A::Nil 
	    graph ++  all//adding all the vertices
     
     
        graph.addEdge(B,A)
        graph.addEdge(F,A)

        
        P(b,"b1",Map[Variable[String],String](),0.005f);
        P(b,"b2",Map[Variable[String],String](),0.995f);
        
        P(f,"f1",Map[Variable[String],String](),0.03f);
        P(f,"f2",Map[Variable[String],String](),0.97f);
        
        P(a,"a1",Map[Variable[String],String]() + ((b,"b1"))+((f,"f1")),0.992f);
        P(a,"a2",Map[Variable[String],String]() + ((b,"b1"))+((f,"f1")),0.008f);
        P(a,"a1",Map[Variable[String],String]() + ((b,"b2"))+((f,"f1")),0.2f);
        P(a,"a2",Map[Variable[String],String]() + ((b,"b2"))+((f,"f1")),0.8f);
        P(a,"a1",Map[Variable[String],String]() + ((b,"b1"))+((f,"f2")),0.99f);
        P(a,"a2",Map[Variable[String],String]() + ((b,"b1"))+((f,"f2")),0.01f);
        P(a,"a1",Map[Variable[String],String]() + ((b,"b2"))+((f,"f2")),0.003f);
        P(a,"a2",Map[Variable[String],String]() + ((b,"b2"))+((f,"f2")),0.997f);
        


}
